from torchvision import transforms
from PIL import ImageEnhance
from . import TRANSFORMS
from .image_jitter import ImageJitter


@TRANSFORMS.register("to_tensor")
def to_tensor(cfg):
    return transforms.ToTensor()


@TRANSFORMS.register("random_resized_crop")
def random_resized_crop(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    # 先随机缩放, 然后再随机裁剪到指定size大小
    return transforms.RandomResizedCrop(
        size=size,
        scale=cfg.TRANSFORM.RANDOM_RESIZED_CROP.SCALE,
        ratio=cfg.TRANSFORM.RANDOM_RESIZED_CROP.RATIO,
    )


@TRANSFORMS.register("random_crop")
def random_crop(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    # 直接随机裁剪到指定size大小
    return transforms.RandomCrop(
        size, padding=cfg.TRANSFORM.RANDOM_CROP.PADDING
    )


@TRANSFORMS.register("random_horizontal_flip")
def random_horizontal_flip(cfg):
    return transforms.RandomHorizontalFlip()


@TRANSFORMS.register("random_vertical_flip")
def random_vertical_flip(cfg):
    return transforms.RandomVerticalFlip()


@TRANSFORMS.register("resize_for_crop")
def resize_for_crop(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    assert size[0] == size[1], "this img-process only process square-image"
    # 把图片短边缩放到指定大小的size/0.875, 然后用别的crop再裁剪到size大小
    return transforms.Resize((int(size[0] / 0.875), int(size[0] / 0.875)))


@TRANSFORMS.register("resize")
def normal_resize(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    return transforms.Resize(size)


@TRANSFORMS.register("center_crop")
def center_crop(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    return transforms.CenterCrop(size)


@TRANSFORMS.register("ten_crop")
def ten_crop(cfg):
    size = cfg.DATASET.IMAGE_SIZE
    # transforms.FiveCrop就是在原图片的四个角和中心各截取一幅大小为size的图片，
    # 而transforms.TenCrop就是在transforms.FiveCrop基础上再进行水平或者竖直翻转（Flip），默认为水平翻转。
    return transforms.TenCrop(size)


@TRANSFORMS.register("normalize")
def normalize(cfg):
    return transforms.Normalize(
        mean=cfg.TRANSFORM.NORMALIZE.MEAN,
        std=cfg.TRANSFORM.NORMALIZE.STD,
    )


@TRANSFORMS.register("random_gray_scale")
def random_gray_scale(cfg):
    # 以一定的概率对图像进行灰度化，转换后的图片还是3通道的
    return transforms.RandomGrayscale(p=0.2)


@TRANSFORMS.register("image_jitter")
def image_jitter(cfg):
    jitter_param = dict(Brightness=cfg.TRANSFORM.IMAGE_JITTER.BRIGHTNESS,
                        Contrast=cfg.TRANSFORM.IMAGE_JITTER.CONTRAST,
                        Color=cfg.TRANSFORM.IMAGE_JITTER.COLOR)
    return ImageJitter(jitter_param)


@TRANSFORMS.register("color_jitter")
def color_jitter(cfg):
    # 颜色扭曲操作（color-jitter），它可以随机调整图像的亮度、对比度、饱和度和色相。
    s = 1

    brightness = cfg.TRANSFORM.COLOR_JITTER.BRIGHTNESS
    contrast = cfg.TRANSFORM.COLOR_JITTER.CONTRAST
    saturation = cfg.TRANSFORM.COLOR_JITTER.SATURATION
    hue = cfg.TRANSFORM.COLOR_JITTER.HUE

    color_jitter = transforms.ColorJitter(brightness * s, contrast * s, saturation * s, hue * s)
    return transforms.RandomApply([color_jitter], p=0.8)


@TRANSFORMS.register('rotate90')
def rotate_90(cfg):
    return transforms.RandomRotation(degrees=(90, 90), expand=False)


@TRANSFORMS.register('rotate180')
def rotate_180(cfg):
    return transforms.RandomRotation(degrees=(180, 180), expand=False)


@TRANSFORMS.register('rotate270')
def rotate_270(cfg):
    return transforms.RandomRotation(degrees=(270, 270), expand=False)
